import os.path

import tqdm
from gxl_ai_utils.utils import utils_file
import librosa


def get_ltr():
    """"""
    tsv_path = "/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/chat_1500"
    txt_path = "/home/work_nfs11/code/xlgeng/txt.lst"
    txt_output_dir = "/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/chat_1500/txt"
    # handle txt
    lines = utils_file.load_list_file_clean(txt_path)
    txt_dict = {}
    for line in lines:
        line = line.split("|")
        txt = line[1]
        txt = txt.replace(" ", "")
        txt = " ".join(txt)
        txt_dict[line[0]] = txt
    # utils_file.print_dict(txt_dict)
    # handle tsv
    partitions = ["test", "valid", "train"]
    for partition in partitions:
        tsv_file = tsv_path + "/" + partition + ".tsv"
        lines = utils_file.load_list_file_clean(tsv_file)
        root_dir = lines[0]
        lines = lines[1:]
        wav_paths = []
        for line in lines:
            line_items = line.split()
            line = line_items[0]
            wav_path = os.path.join(root_dir, line)
            wav_paths.append(wav_path)
        wav_dict = {}
        for wav_path in wav_paths:
            wav_name = utils_file.get_file_pure_name_from_path(wav_path)
            wav_dict[wav_name] = wav_path
        # utils_file.print_dict(wav_dict)
        res_txt =  []
        for key in wav_dict.keys():
            txt = txt_dict[key]
            res_txt.append(txt)
        # utils_file.print_list(res_txt)
        ltr_path = txt_output_dir + "/" + partition + ".ltr"
        utils_file.write_list_to_file(res_txt, ltr_path)
        dict_set = set()
        if partition == "train":
            # 得到dict
            for txt in res_txt:
                chars = txt.split(" ")
                for char in chars:
                    dict_set.add(char)
            res_dict_dict = {}
            for i, char in enumerate(dict_set):
                res_dict_dict[char] = i
            utils_file.print_dict(res_dict_dict)
            print(dict_set)
            output_path = txt_output_dir + "/" + "dict.ltr.txt"
            utils_file.write_dict_to_scp(res_dict_dict, output_path)


def get_sample_rate(file_path):
    samplerate = librosa.get_samplerate(file_path)
    return samplerate
def filter_data():
    tsv_path = "/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/chat_1500"
    outtput_tsv_path = "/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/chat_1500/new2"
    utils_file.makedir(outtput_tsv_path)
    partions = ['train','dev','test']
    for partition in partions:
        tsv_file = tsv_path + "/" + partition + ".tsv"
        ltr_file = tsv_path + "/" + partition + ".ltr"
        tsv_lines = utils_file.load_list_file_clean(tsv_file)
        ltr_lines = utils_file.load_list_file_clean(ltr_file)
        root_dir = tsv_lines[0]
        res_tsv_lines = tsv_lines[1:]
        assert len(res_tsv_lines)==len(ltr_lines), "tsv and ltr files are not of the same length!"
        true_tsv_lines = [root_dir]
        true_ltr_lines  = []
        index_true = 0
        index_false = 0
        for wav_path_with_sample, text in tqdm.tqdm(zip(res_tsv_lines, ltr_lines), total=len(res_tsv_lines)):
            wav_path = wav_path_with_sample.split()[0]
            wav_path = os.path.join(root_dir, wav_path)
            sample_rate = get_sample_rate(wav_path)
            if sample_rate!=16000:
                index_false += 1
                if index_false % 100 == 0:
                    print(index_false)
                continue
            index_true += 1
            true_tsv_lines.append(wav_path_with_sample)
            true_ltr_lines.append(text)
            if index_true % 10000==0:
                output_tsv_path = outtput_tsv_path + "/" + partition + ".tsv"
                utils_file.write_list_to_file(true_tsv_lines, output_tsv_path)
                output_ltr_path = outtput_tsv_path + "/" + partition + ".ltr"
                utils_file.write_list_to_file(true_ltr_lines, output_ltr_path)
        output_tsv_path = outtput_tsv_path + "/" + partition + ".tsv"
        utils_file.write_list_to_file(true_tsv_lines, output_tsv_path)
        output_ltr_path = outtput_tsv_path + "/" + partition + ".ltr"
        utils_file.write_list_to_file(true_ltr_lines, output_ltr_path)














if __name__ == "__main__":
    filter_data()